from tensorflow.keras.metrics import MeanIoU
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import confusion_matrix
from tensorflow.python.framework import dtypes


class VeinIou(MeanIoU):
    def __init__(self, num_classes, name=None, dtype=None):
        super(VeinIou, self).__init__(num_classes, name=None, dtype=None)

    def update_state(self, y_true, y_pred, sample_weight=None):
        y_true = math_ops.cast(y_true[:, :, :0:self.num_classes], self._dtype)
        y_pred = math_ops.cast(y_pred[:, :, :0:self.num_classes], self._dtype)

        # Flatten the input if its rank > 1.
        if y_pred.shape.ndims > 1:
            y_pred = array_ops.reshape(y_pred, [-1])

        if y_true.shape.ndims > 1:
            y_true = array_ops.reshape(y_true, [-1])

        if sample_weight is not None:
            sample_weight = math_ops.cast(sample_weight, self._dtype)
            if sample_weight.shape.ndims > 1:
                sample_weight = array_ops.reshape(sample_weight, [-1])

        # Accumulate the prediction to current confusion matrix.
        current_cm = confusion_matrix.confusion_matrix(
            y_true,
            y_pred,
            self.num_classes,
            weights=sample_weight,
            dtype=dtypes.float64)
        return self.total_cm.assign_add(current_cm)

